#!/usr/bin/env python3
import rospy
from sensor_msgs.msg import Joy
import command
import time
import LED
import config
import Angle
from Logger import LogToTopic

# 控制周期（秒），50Hz刷新频率
CONTROL_PERIOD = 0.04
# 摇杆死区阈值
DEADZONE_THRESHOLD = 0.1

class JoystickDemo:
    def __init__(self):
        rospy.init_node('joystick_demo')
        
        # 初始化组件
        self.logger = LogToTopic('/ros_log')
        self.led = LED.LED()
        self.robot = command.Robot()
        
        # 配置参数
        self.snake_speed = config.SNAKE_SPEED
        self.speed_threshold = config.SPEED_THRESHOLD
        self.ANGLE = config.ANGLE_THRESHOLD
        
        # 状态变量初始化
        self.axes = [0.0] * 8
        self.buttons = [0] * 12
        
        self.LED_state = False  # False表示关，True表示开
        self.Tail_state = False
        self.Mode_state = False  # False代表速度模式（默认），True代表位置模式
        
        # 关节角度存储
        self.joint_angles = {
            2: 0,
            3: 0,
            4: 0,
            5: 0
        }
        
        # 上一周期状态记录
        self.last_speed = None
        self.last_turn_state = 0  # -1:左, 0:停止, 1:右
        self.last_head_state = 0  # -1:低头, 0:停止, 1:抬头
        self.last_button_states = [0] * 12
        self.last_axis_states = [0.0] * 8
        
        # 订阅手柄话题
        rospy.Subscriber("joy", Joy, self.joy_callback)
        rospy.on_shutdown(self.shutdown)
        
        # 启动主控制循环
        self.main_loop()

    def joy_callback(self, data):
        """处理手柄原始数据"""
        self.axes = data.axes
        self.buttons = data.buttons

    def get_joint_angles(self):
        """获取关节角度数据"""
        for joint in [2, 3, 4, 5]:
            angle = self.robot.get_angle_data(joint)
            if angle is not None:
                # 根据关节限制过滤角度值
                if joint in [2, 5] and abs(angle) < 45:
                    self.joint_angles[joint] = angle
                elif joint in [3, 4] and abs(angle) < 85:
                    self.joint_angles[joint] = angle

    def process_movement(self):
        """处理移动控制"""
        # 前进后退控制
        axis_value = self.axes[1]
        if abs(axis_value) > DEADZONE_THRESHOLD:
            raw_speed = axis_value * self.snake_speed
            target_speed = int(round(raw_speed / self.speed_threshold) * self.speed_threshold)
            
            if target_speed != self.last_speed:
                if target_speed > 0:
                    rospy.loginfo(f"履带蛇前进，速度: {target_speed}")
                    self.robot.move_front(motor_name=1, speed=target_speed)
                elif target_speed < 0:
                    rospy.loginfo(f"履带蛇后退，速度: {target_speed}")
                    self.robot.move_front(motor_name=1, speed=target_speed)
                self.last_speed = target_speed
        else:
            if self.last_speed is not None:
                rospy.loginfo("履带蛇停止运动")
                self.robot.stop_move()
                self.last_speed = None

        # 左右转向控制
        turn_value = self.axes[2]
        current_turn_state = 0
        if turn_value > DEADZONE_THRESHOLD:
            current_turn_state = 1  # 右转
        elif turn_value < -DEADZONE_THRESHOLD:
            current_turn_state = -1  # 左转
            
        if current_turn_state != self.last_turn_state:
            if current_turn_state == 1:
                rospy.loginfo("履带蛇右转")
                self.robot.move_joint_speed(5, -1)
                self.robot.move_joint_speed(2, 1)
            elif current_turn_state == -1:
                rospy.loginfo("履带蛇左转")
                self.robot.move_joint_speed(5, 1)
                self.robot.move_joint_speed(2, -1)
            else:
                rospy.loginfo("履带蛇停止左右转")
                self.robot.move_joint_speed(5, 0)
                self.robot.move_joint_speed(2, 0)
            self.last_turn_state = current_turn_state

        # 头部控制
        head_up = self.buttons[3] == 1
        head_down = self.buttons[1] == 1
        current_head_state = 0
        if head_up:
            current_head_state = 1
        elif head_down:
            current_head_state = -1
            
        if current_head_state != self.last_head_state:
            if current_head_state == 1:
                rospy.loginfo("履带蛇抬头")
                self.robot.move_joint_speed(3, 1)
            elif current_head_state == -1:
                rospy.loginfo("履带蛇低头")
                self.robot.move_joint_speed(3, -1)
            else:
                rospy.loginfo("停止头部运动")
                self.robot.move_joint_speed(3, 0)
            self.last_head_state = current_head_state

    def process_buttons(self):
        """处理按钮事件"""
        # 处理使能/失能按钮（8和9）
        if self.buttons[8] == 1 and self.last_button_states[8] == 0:
            rospy.loginfo("电机失能")
            self.robot.disable_all()
        
        if self.buttons[9] == 1 and self.last_button_states[9] == 0:
            rospy.loginfo("电机使能")
            self.robot.enable_all()
        
        # 一键拉直（LB按键4）
        if self.buttons[4] == 1 and self.last_button_states[4] == 0:
            rospy.loginfo("一键拉直")
            self.robot.move_straight1()
        
        # 尾门控制（RB按键5）
        if self.buttons[5] == 1 and self.last_button_states[5] == 0:
            if not self.Tail_state:
                rospy.loginfo("尾门关闭")
            else:
                rospy.loginfo("尾门开启")
            self.Tail_state = not self.Tail_state
        
        # LED控制（RT按键7）
        if self.buttons[7] == 1 and self.last_button_states[7] == 0:
            rospy.loginfo("LED控制")
            if not self.LED_state:
                self.led.off()
                rospy.loginfo("LED关灯")
            else:
                self.led.on()
                rospy.loginfo("LED开灯")
            self.LED_state = not self.LED_state
        
        # 模式切换（LT按键6）
        if self.buttons[6] == 1 and self.last_button_states[6] == 0:
            if not self.Mode_state:
                self.robot.Position_mode()
                rospy.loginfo("切换位置模式")
                # 位置模式下微调关节
                self.robot.move_joint_angle(2, self.joint_angles[2])
                self.robot.move_joint_angle(3, self.joint_angles[3])
                self.robot.move_joint_angle(4, self.joint_angles[4])
                self.robot.move_joint_angle(5, self.joint_angles[5])
            else:
                self.robot.Velocity_mode()
                rospy.loginfo("切换速度模式")
            self.Mode_state = not self.Mode_state
        
        # 记录当前按钮状态作为下次参考
        self.last_button_states = self.buttons.copy()

    def process_axis_buttons(self):
        """处理摇杆轴按键事件（方向键）"""
        # 抬头（轴5向上）
        if self.axes[5] == 1.0 and self.last_axis_states[5] != 1.0:
            rospy.loginfo(f"抬头{self.ANGLE}度")
            if abs(self.joint_angles[3]) <= 80:
                self.joint_angles[3] += self.ANGLE
                self.joint_angles[4] += self.ANGLE
                self.robot.move_joint_angle(3, self.joint_angles[3])
                self.robot.move_joint_angle(4, self.joint_angles[4])
        
        # 低头（轴5向下）
        if self.axes[5] == -1.0 and self.last_axis_states[5] != -1.0:
            rospy.loginfo(f"低头{self.ANGLE}度")
            if abs(self.joint_angles[3]) <= 80:
                self.joint_angles[3] -= self.ANGLE
                self.joint_angles[4] -= self.ANGLE
                self.robot.move_joint_angle(3, self.joint_angles[3])
                self.robot.move_joint_angle(4, self.joint_angles[4])
        
        # 左转（轴4向左）
        if self.axes[4] == 1.0 and self.last_axis_states[4] != 1.0:
            rospy.loginfo(f"左转{self.ANGLE}度")
            if abs(self.joint_angles[2]) <= 35:
                self.joint_angles[2] -= self.ANGLE
                self.joint_angles[5] -= self.ANGLE
                self.robot.move_joint_angle(2, self.joint_angles[2])
                self.robot.move_joint_angle(5, self.joint_angles[5])
        
        # 右转（轴4向右）
        if self.axes[4] == -1.0 and self.last_axis_states[4] != -1.0:
            rospy.loginfo(f"右转{self.ANGLE}度")
            if abs(self.joint_angles[2]) <= 35:
                self.joint_angles[2] += self.ANGLE
                self.joint_angles[5] += self.ANGLE
                self.robot.move_joint_angle(2, self.joint_angles[2])
                self.robot.move_joint_angle(5, self.joint_angles[5])
        
        # 记录当前轴状态作为下次参考
        self.last_axis_states = self.axes.copy()

    def main_loop(self):
        """主控制循环"""
        rate = rospy.Rate(1 / CONTROL_PERIOD)  # 50Hz
        angle_update_counter = 0  # 角度更新计数器
        
        while not rospy.is_shutdown():
            # 每50个周期（约1秒）更新一次关节角度
            if angle_update_counter >= 50:
                self.get_joint_angles()
                angle_update_counter = 0
            else:
                angle_update_counter += 1
            
            # 处理各类控制
            self.process_movement()
            self.process_buttons()
            self.process_axis_buttons()
            
            rate.sleep()

    def shutdown(self):
        """关闭时的清理操作"""
        rospy.loginfo("正在关闭节点...")
        self.robot.stop_move()
        self.robot.disable_all()
        self.logger.shutdown()
        rospy.loginfo("节点已关闭")

if __name__ == '__main__':
    try:
        demo = JoystickDemo()
    except rospy.ROSInterruptException:
        pass